import torch
import torch.nn as nn
from torch_geometric.nn import GCNConv

class Encoder(nn.Module):
    """
    DGI模型的编码器,使用GCN从节点特征和图结构中学习节点的Embedding
    """
    def __init__(self, in_channels, hidden_channels):
        """
        初始化编码器。
        :param in_channels: int, 输入节点特征的维度。
        :param hidden_channels: int, 输出节点嵌入的维度。
        """
        super(Encoder, self).__init__()
        # 使用 PyTorch Geometric 的 GCNConv 层
        self.conv = GCNConv(in_channels, hidden_channels)
        self.prelu = nn.PReLU(hidden_channels)

    def forward(self, x, edge_index):
        x = self.conv(x, edge_index)
        x = self.prelu(x)
        return x

class Discriminator(nn.Module):
    # DGI模型的判别器，用于区分一个节点表示（局部）是否来自于一个给定的图摘要（全局）
    def __init__(self, in_channels):
        super(Discriminator, self).__init__()
        # 双线性层，用于计算节点表示和图摘要之间的分数
        self.bilinear = nn.Bilinear(in_channels, in_channels, 1)

    def forward(self, node_embedding, graph_summary):
        # 为了匹配双线性层的输入要求，需要将图摘要向量扩展
        graph_summary_expanded = graph_summary.expand_as(node_embedding)
        
        scores = self.bilinear(node_embedding, graph_summary_expanded)
        
        return scores

class DGI(nn.Module):
    def __init__(self, in_channels, hidden_channels):

        super(DGI, self).__init__()
        self.encoder = Encoder(in_channels, hidden_channels)
        self.discriminator = Discriminator(hidden_channels)
        self.loss_fn = nn.BCEWithLogitsLoss()

    def forward(self, data):
        """
        DGI模型的前向传播过程。
        :param data: PyG的Data对象，或任何包含'x'和'edge_index'的对象/字典。
        :return: tuple, (positive_scores, negative_scores)
                 positive_scores: 真实节点表示与图摘要之间的分数。
                 negative_scores: 损坏节点表示与图摘要之间的分数。
        """
        features, edge_index = data.x, data.edge_index

        # 1. 获取正样本的节点嵌入 H
        positive_embeds = self.encoder(features, edge_index)
        
        # 2. 生成负样本的节点特征（通过行置换/shuffling），腐化函数
        shuffled_idx = torch.randperm(features.shape[0])
        corrupted_features = features[shuffled_idx]

        # 3. 获取负样本的节点嵌入 H_tilde
        negative_embeds = self.encoder(corrupted_features, edge_index)

        # 4. 计算图的全局摘要向量 s (Readout function R)
        summary_vector = torch.sigmoid(positive_embeds.mean(dim=0))

        # 5. 使用判别器 D 进行评分
        positive_scores = self.discriminator(positive_embeds, summary_vector)
        negative_scores = self.discriminator(negative_embeds, summary_vector)

        return positive_scores, negative_scores

    def loss(self, pos_scores, neg_scores):
        """
        计算DGI的损失函数。
        目标是让正样本对的分数接近1，负样本对的分数接近0。
        :param pos_scores: Tensor, 正样本对的分数。
        :param neg_scores: Tensor, 负样本对的分数。
        :return: Tensor, 计算出的总损失。
        """
        # 生成标签：正样本为1，负样本为0
        labels = torch.cat([
            torch.ones_like(pos_scores),
            torch.zeros_like(neg_scores)
        ])

        scores = torch.cat([pos_scores, neg_scores])
        
        return self.loss_fn(scores, labels)

    def get_embeds(self, data):
        # 在模型训练完成后，用此方法获取节点的最终嵌入。
        self.eval()
        with torch.no_grad():
            embeds = self.encoder(data.x, data.edge_index)
        return embeds